
package org.example.langchan4j.config;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowMapper;

import java.sql.Timestamp;
import java.time.LocalDateTime;
import java.util.List;

public class PersistentChatMemoryStore implements ChatMemoryStore {

    private final JdbcTemplate jdbcTemplate;
    private final ObjectMapper objectMapper;

    public PersistentChatMemoryStore(JdbcTemplate jdbcTemplate) {
        this.jdbcTemplate = jdbcTemplate;
        this.objectMapper = new ObjectMapper();
    }

    private static final RowMapper<List<ChatMessage>> CHAT_MESSAGE_ROW_MAPPER = (rs, rowNum) -> {
        String messagesJson = rs.getString("messages");
        ObjectMapper mapper = new ObjectMapper();
        try {
            return mapper.readValue(messagesJson, new TypeReference<List<ChatMessage>>() {});
        } catch (Exception e) {
            throw new RuntimeException("Failed to deserialize chat messages", e);
        }
    };

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        String chatId = memoryId.toString();
        String sql = "SELECT messages FROM chat_memory WHERE chat_id = ? ORDER BY updated_at DESC LIMIT 1";

        List<List<ChatMessage>> result = jdbcTemplate.query(sql, CHAT_MESSAGE_ROW_MAPPER, chatId);
        if (!result.isEmpty()) {
            return result.get(0);
        }
        return List.of();
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        String chatId = memoryId.toString();
        try {
            String messagesJson = objectMapper.writeValueAsString(messages);

            // 检查记录是否存在
            String checkSql = "SELECT COUNT(*) FROM chat_memory WHERE chat_id = ?";
            Integer count = jdbcTemplate.queryForObject(checkSql, Integer.class, chatId);

            if (count != null && count > 0) {
                // 更新现有记录
                String updateSql = "UPDATE chat_memory SET messages = ?, updated_at = ? WHERE chat_id = ?";
                jdbcTemplate.update(updateSql, messagesJson, Timestamp.valueOf(LocalDateTime.now()), chatId);
            } else {
                // 插入新记录
                String insertSql = "INSERT INTO chat_memory (chat_id, messages, created_at, updated_at) VALUES (?, ?, ?, ?)";
                jdbcTemplate.update(insertSql, chatId, messagesJson,
                        Timestamp.valueOf(LocalDateTime.now()), Timestamp.valueOf(LocalDateTime.now()));
            }
        } catch (Exception e) {
            throw new RuntimeException("Failed to serialize and save chat messages", e);
        }
    }

    @Override
    public void deleteMessages(Object memoryId) {
        String chatId = memoryId.toString();
        String sql = "DELETE FROM chat_memory WHERE chat_id = ?";
        jdbcTemplate.update(sql, chatId);
    }
}